package com.apisix.service.controller;

import com.apisix.service.entity.Consumer;
import com.apisix.service.service.ConsumerService;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.SignatureAlgorithm;
import io.jsonwebtoken.security.Keys;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestHeader;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.util.Base64;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

@RestController
@RequestMapping("/api/token")
public class TokenController {

    private final ConsumerService consumerService;
    
    // 时间戳有效期（毫秒）
    private static final long TIMESTAMP_VALIDITY = 5 * 60 * 1000; // 5分钟

    @Autowired
    public TokenController(ConsumerService consumerService) {
        this.consumerService = consumerService;
    }

    /**
     * 获取JWT令牌
     * 
     * @param authorization 授权签名，是app_secret+时间戳的MD5值（大写）
     * @param accessKey 消费者的app_id
     * @param timestamp 时间戳，用于防止重放攻击
     * @return JWT令牌
     */
    @GetMapping
    public ResponseEntity<?> getToken(
            @RequestHeader("Authorization") String authorization,
            @RequestHeader("Accesskey") String accessKey,
            @RequestParam("t") String timestamp) {
        
        try {
            // 验证时间戳是否在有效期内
            long currentTime = System.currentTimeMillis();
            long requestTime;
            try {
                requestTime = Long.parseLong(timestamp);
            } catch (NumberFormatException e) {
                return ResponseEntity.status(HttpStatus.BAD_REQUEST)
                        .body(Map.of("error", "无效的时间戳格式"));
            }
            
            // 检查时间戳是否过期
            if (Math.abs(currentTime - requestTime) > TIMESTAMP_VALIDITY) {
                return ResponseEntity.status(HttpStatus.UNAUTHORIZED)
                        .body(Map.of("error", "时间戳已过期或无效"));
            }
            
            // 根据accessKey获取消费者信息
            Optional<Consumer> consumerOpt = consumerService.getConsumerByAppId(accessKey);
            
            if (consumerOpt.isEmpty()) {
                return ResponseEntity.status(HttpStatus.UNAUTHORIZED)
                        .body(Map.of("error", "无效的AccessKey"));
            }
            
            Consumer consumer = consumerOpt.get();
            
            // 验证授权签名 - 使用app_secret+时间戳生成MD5
            String expectedAuthorization = getMD5Hash(consumer.getAppSecret() + timestamp);
            
            if (!authorization.equalsIgnoreCase(expectedAuthorization)) {
                return ResponseEntity.status(HttpStatus.UNAUTHORIZED)
                        .body(Map.of("error", "授权签名无效"));
            }
            
            // 生成JWT令牌
            String token = generateJwtToken(consumer);
            
            Map<String, Object> response = new HashMap<>();
            response.put("token", token);
            response.put("expires_in", 3600); // 令牌有效期，单位秒
            
            return ResponseEntity.ok(response);
            
        } catch (Exception e) {
            return ResponseEntity.status(HttpStatus.INTERNAL_SERVER_ERROR)
                    .body(Map.of("error", "生成令牌失败: " + e.getMessage()));
        }
    }
    
    /**
     * 生成JWT令牌
     * 
     * @param consumer 消费者信息
     * @return JWT令牌字符串
     */
    private String generateJwtToken(Consumer consumer) throws Exception {
        long now = System.currentTimeMillis();
        Date issuedAt = new Date(now);
        Date expiresAt = new Date(now + 3600 * 1000); // 1小时有效期
        
        Map<String, Object> claims = new HashMap<>();
        claims.put("key", consumer.getAppId());
        
        // 如果存在RSA私钥，使用RSA算法签名
        if (consumer.getRsaPrivateKey() != null && !consumer.getRsaPrivateKey().isEmpty()) {
            PrivateKey privateKey = getPrivateKeyFromString(consumer.getRsaPrivateKey());
            
            return Jwts.builder()
                    .setClaims(claims)
                    .setIssuedAt(issuedAt)
                    .setExpiration(expiresAt)
                    .signWith(privateKey, SignatureAlgorithm.RS256)
                    .compact();
        } else {
            // 否则使用HMAC算法签名
            return Jwts.builder()
                    .setClaims(claims)
                    .setIssuedAt(issuedAt)
                    .setExpiration(expiresAt)
                    .signWith(Keys.hmacShaKeyFor(consumer.getAppSecret().getBytes(StandardCharsets.UTF_8)), 
                            SignatureAlgorithm.HS256)
                    .compact();
        }
    }
    
    /**
     * 从PEM格式字符串中获取私钥
     * 
     * @param privateKeyPEM PEM格式的私钥字符串
     * @return PrivateKey对象
     */
    private PrivateKey getPrivateKeyFromString(String privateKeyPEM) throws Exception {
        String privateKeyContent = privateKeyPEM
                .replace("-----BEGIN PRIVATE KEY-----", "")
                .replace("-----END PRIVATE KEY-----", "")
                .replaceAll("\\s", "");
        
        byte[] privateKeyBytes = Base64.getDecoder().decode(privateKeyContent);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        
        return keyFactory.generatePrivate(keySpec);
    }
    
    /**
     * 计算字符串的MD5哈希值（大写）
     * 
     * @param input 输入字符串
     * @return MD5哈希值（大写）
     */
    private String getMD5Hash(String input) throws NoSuchAlgorithmException {
        MessageDigest md = MessageDigest.getInstance("MD5");
        byte[] hashBytes = md.digest(input.getBytes(StandardCharsets.UTF_8));
        
        StringBuilder sb = new StringBuilder();
        for (byte b : hashBytes) {
            sb.append(String.format("%02X", b));
        }
        
        return sb.toString();
    }
} 